{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "9c48213d-6e6a-4c10-838a-2a7c710c3a05",
   "metadata": {},
   "source": [
    "# Guardrails Output Parsing"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "50d3b817-b70e-4667-be4f-d3a0fe4bd119",
   "metadata": {},
   "source": [
    "#### Load documents, build the VectorStoreIndex"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "690a6918-7c75-4f95-9ccc-d2c4a1fe00d7",
   "metadata": {},
   "outputs": [],
   "source": [
    "import logging\n",
    "import sys\n",
    "\n",
    "logging.basicConfig(stream=sys.stdout, level=logging.INFO)\n",
    "logging.getLogger().addHandler(logging.StreamHandler(stream=sys.stdout))\n",
    "\n",
    "from llama_index import VectorStoreIndex, SimpleDirectoryReader\n",
    "from IPython.display import Markdown, display"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "03d1691e-544b-454f-825b-5ee12f7faa8a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# load documents\n",
    "documents = SimpleDirectoryReader(\"../paul_graham_essay/data\").load_data()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "ad144ee7-96da-4dd6-be00-fd6cf0c78e58",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:llama_index.token_counter.token_counter:> [build_index_from_documents] Total LLM token usage: 0 tokens\n",
      "> [build_index_from_documents] Total LLM token usage: 0 tokens\n",
      "INFO:llama_index.token_counter.token_counter:> [build_index_from_documents] Total embedding token usage: 18579 tokens\n",
      "> [build_index_from_documents] Total embedding token usage: 18579 tokens\n"
     ]
    }
   ],
   "source": [
    "index = VectorStoreIndex.from_documents(documents, chunk_size=512)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8b7d7c61-b5d7-4b8f-b90b-3ebee1103f27",
   "metadata": {},
   "source": [
    "#### Define Query + Guardrails Spec"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "6fb88295-0840-4e2d-b79b-def0b0a63a7f",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from llama_index.output_parsers import GuardrailsOutputParser\n",
    "from llama_index.llm_predictor import StructuredLLMPredictor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "057139d2-09e8-4b8d-83a1-a2356a1475a8",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "llm_predictor = StructuredLLMPredictor()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bc25edf7-9343-4e82-a3f1-eec4281a9371",
   "metadata": {},
   "source": [
    "**Define custom QA and Refine Prompts**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "2833d086-d240-4798-b3c5-a83ac4593b0e",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from llama_index.prompts import PromptTemplate\n",
    "from llama_index.prompts.default_prompts import (\n",
    "    DEFAULT_TEXT_QA_PROMPT_TMPL,\n",
    "    DEFAULT_REFINE_PROMPT_TMPL,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "a4b9201d-fe16-4cc0-8135-a08d9928625d",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# NOTE: we don't need to define the query_str in the rail spec, we can define during query-time.\n",
    "rail_spec = \"\"\"\n",
    "<rail version=\"0.1\">\n",
    "\n",
    "<output>\n",
    "    <list name=\"points\" description=\"Bullet points regarding events in the author's life.\">\n",
    "        <object>\n",
    "            <string name=\"explanation\" format=\"one-line\" on-fail-one-line=\"noop\" />\n",
    "            <string name=\"explanation2\" format=\"one-line\" on-fail-one-line=\"noop\" />\n",
    "            <string name=\"explanation3\" format=\"one-line\" on-fail-one-line=\"noop\" />\n",
    "        </object>\n",
    "    </list>\n",
    "</output>\n",
    "\n",
    "<prompt>\n",
    "\n",
    "Query string here.\n",
    "\n",
    "@xml_prefix_prompt\n",
    "\n",
    "{output_schema}\n",
    "\n",
    "@json_suffix_prompt_v2_wo_none\n",
    "</prompt>\n",
    "</rail>\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "f7af4ebf-1dff-48ec-9fb7-8926af45b6a0",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "output_parser = GuardrailsOutputParser.from_rail_string(\n",
    "    rail_spec, llm=llm_predictor.llm\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "a9b440d4-6fb4-46e6-973f-44207b432d3f",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# NOTE: we use the same output parser for both prompts, though you can choose to use different parsers\n",
    "# NOTE: here we add formatting instructions to the prompts.\n",
    "\n",
    "fmt_qa_tmpl = output_parser.format(DEFAULT_TEXT_QA_PROMPT_TMPL)\n",
    "fmt_refine_tmpl = output_parser.format(DEFAULT_REFINE_PROMPT_TMPL)\n",
    "\n",
    "qa_prompt = PromptTemplate(fmt_qa_tmpl, output_parser=output_parser)\n",
    "refine_prompt = PromptTemplate(fmt_refine_tmpl, output_parser=output_parser)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "1ba18a80-35f4-4fd4-9b13-9f13f84db4fe",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Context information is below. \n",
      "---------------------\n",
      "{context_str}\n",
      "---------------------\n",
      "Given the context information and not prior knowledge, answer the question: {query_str}\n",
      "\n",
      "\n",
      "\n",
      "Given below is XML that describes the information to extract from this document and the tags to extract it into.\n",
      "\n",
      "\n",
      "<output>\n",
      "    <list name=\"points\" description=\"Bullet points regarding events in the author's life.\">\n",
      "        <object>\n",
      "            <string name=\"explanation\" format=\"one-line\"/>\n",
      "            <string name=\"explanation2\" format=\"one-line\"/>\n",
      "            <string name=\"explanation3\" format=\"one-line\"/>\n",
      "        </object>\n",
      "    </list>\n",
      "</output>\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "ONLY return a valid JSON object (no other text is necessary). The JSON MUST conform to the XML format, including any types and format requests e.g. requests for lists, objects and specific types. Be correct and concise.\n",
      "\n",
      "JSON Output:\n",
      "\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# take a look at the new QA template!\n",
    "print(fmt_qa_tmpl)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b6caf93b-6345-4c65-a346-a95b0f1746c4",
   "metadata": {},
   "source": [
    "#### Query Index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "fb9cdf43-0f31-4c36-869b-df9fa50aebdb",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:llama_index.token_counter.token_counter:> [query] Total LLM token usage: 754 tokens\n",
      "> [query] Total LLM token usage: 754 tokens\n",
      "INFO:llama_index.token_counter.token_counter:> [query] Total embedding token usage: 11 tokens\n",
      "> [query] Total embedding token usage: 11 tokens\n"
     ]
    }
   ],
   "source": [
    "query_engine = index.as_query_engine(\n",
    "    text_qa_template=qa_prompt,\n",
    "    refine_template=refine_prompt,\n",
    "    llm_predictor=llm_predictor,\n",
    ")\n",
    "response = query_engine.query(\n",
    "    \"What are the three items the author did growing up?\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "bc7760b6-5be3-4303-b97e-3f5edacf674b",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'points': [{'explanation': 'Writing short stories', 'explanation2': 'Programming on an IBM 1401', 'explanation3': 'Using microcomputers'}]}\n"
     ]
    }
   ],
   "source": [
    "print(response)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
